#ifndef __TILING_COMMON_H_
#define __TILING_COMMON_H_
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"
#include <csignal>
#include <cstdint>
#include <functional>
#include <iostream>
#include <tuple>
#define COMMON_TILING_FILED_DEF                                                \
  TILING_DATA_FIELD_DEF(int64_t, size);                                        \
  TILING_DATA_FIELD_DEF(int64_t, formerNum);                                   \
  TILING_DATA_FIELD_DEF(int64_t, formerLength);                                \
  TILING_DATA_FIELD_DEF(int64_t, formerTileLength);                            \
  TILING_DATA_FIELD_DEF(int64_t, formerFinalCalcCount);                        \
  TILING_DATA_FIELD_DEF(int64_t, tailNum);                                     \
  TILING_DATA_FIELD_DEF(int64_t, tailLength);                                  \
  TILING_DATA_FIELD_DEF(int64_t, tailTileLength);                              \
  TILING_DATA_FIELD_DEF(int64_t, tailFinalCalcCount);                          \
  TILING_DATA_FIELD_DEF(int64_t, finalKernelCaclCount);                        \
  TILING_DATA_FIELD_DEF(int64_t, dType);

#define COMMON_TILING_TAG_DEF TILING_DATA_FIELD_DEF(uint64_t, tag)

#define COMMON_TILING_GET_PLATFORM                                             \
  auto ascendcPlatform =                                                       \
      platform_ascendc::PlatformAscendC(context->GetPlatformInfo());

#define COMMON_TILING_GET_UB_SIZE                                              \
  uint64_t ub_size;                                                            \
  ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ub_size);

#define COMMON_TILING_GET_DTYPE_DATA(input_idx)                                \
  auto dtype = context->GetInputDesc(input_idx)->GetDataType();                \
  auto dtypeSize = GetSizeByDataType(dtype);                                   \
  tiling.set_dType(dtype);

#define COMMON_TILING_ALIGN_DATA(align_size, data_size)                        \
  auto alignDataSize = align_size / dtypeSize;                                 \
  auto alignCount = (data_size + alignDataSize - 1) / alignDataSize;           \
  auto alignedSize = alignCount * alignDataSize;

#define COMMON_TILING_SPLIT(tiling)                                            \
  auto kernelBlock = alignCount / blockDim;                                    \
  if (kernelBlock == 0) {                                                      \
    blockDim = alignCount;                                                     \
    context->SetBlockDim(blockDim);                                            \
    kernelBlock = 1;                                                           \
  }                                                                            \
  auto formerNum = alignCount % blockDim;                                      \
  auto formerLength = (kernelBlock + 1) * alignDataSize;                       \
  tiling.set_formerNum(formerNum);                                             \
  tiling.set_formerLength(formerLength);                                       \
  auto tailNum = blockDim - formerNum;                                         \
  auto tailLength = kernelBlock * alignDataSize;                               \
  tiling.set_tailNum(tailNum);                                                 \
  tiling.set_tailLength(tailLength);

#define CEIL_DIV(a, b) ((a + b - 1) / b)
#define ALIGN_TO(a, b) CEIL_DIV(a, b) * b

namespace optiling {
/*
返回:
alignedSize,
alignDataSize,
blockDim,
formerNum,
formerLength,
tailNum,
tailLength
*/
inline std::tuple<int64_t, int64_t, int64_t, int64_t, int64_t, int64_t, int64_t>
common_tiling_split(int64_t blockDim, int64_t dtypeSize, int64_t alignTo,
                    int64_t dataSize) {
  COMMON_TILING_ALIGN_DATA(alignTo, dataSize);
  auto kernelBlock = alignCount / blockDim;
  if (kernelBlock == 0) {
    blockDim = alignCount;
    kernelBlock = 1;
  }
  auto formerNum = alignCount % blockDim;
  auto formerLength = (kernelBlock + 1) * alignDataSize;
  auto tailNum = blockDim - formerNum;
  auto tailLength = kernelBlock * alignDataSize;
  return std::make_tuple(alignedSize, alignDataSize, blockDim, formerNum,
                         formerLength, tailNum, tailLength);
}
typedef int64_t (*mem_addition_fn)(ge::DataType);
template <typename _T_TILING> void print_common_tiling(_T_TILING &tiling) {
  std::cout << "get_size : " << tiling.get_size() << std::endl;
  std::cout << "get_formerNum : " << tiling.get_formerNum() << std::endl;
  std::cout << "get_formerLength : " << tiling.get_formerLength() << std::endl;
  std::cout << "get_formerTileLength : " << tiling.get_formerTileLength()
            << std::endl;
  std::cout << "get_formerFinalCalcCount : "
            << tiling.get_formerFinalCalcCount() << std::endl;
  std::cout << "get_tailNum : " << tiling.get_tailNum() << std::endl;
  std::cout << "get_tailLength : " << tiling.get_tailLength() << std::endl;
  std::cout << "get_tailTileLength : " << tiling.get_tailTileLength()
            << std::endl;
  std::cout << "get_tailFinalCalcCount : " << tiling.get_tailFinalCalcCount()
            << std::endl;
  std::cout << "get_finalKernelCaclCount : "
            << tiling.get_finalKernelCaclCount() << std::endl;
  std::cout << "get_dType : " << tiling.get_dType() << std::endl;
}

template <typename _T_TILING>
inline void tiling_save_to_buffer(_T_TILING &tiling,
                                  gert::TilingContext *context) {
  tiling.SaveToBuffer(context->GetRawTilingData()->GetData(),
                      context->GetRawTilingData()->GetCapacity());
  context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
}

template <typename _T_TILING>
inline int64_t set_size_by_input(_T_TILING &tiling,
                                 gert::TilingContext *context,
                                 int64_t shape_index) {
  auto shape = context->GetInputShape(shape_index);
  auto dataSize = shape->GetStorageShape().GetShapeSize();
  tiling.set_size(dataSize);
  return dataSize;
}

inline void tiling_zero_workspace(gert::TilingContext *context) {
  size_t *currentWorkspace = context->GetWorkspaceSizes(1);
  currentWorkspace[0] = 0;
}

template <typename _T_TILING>
int64_t common_tiling_set_kernel_vec(_T_TILING &tiling,
                                     gert::TilingContext *context) {
  auto ascendcPlatform =
      platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
  auto aicNum = ascendcPlatform.GetCoreNumAic();
  auto aivNum = ascendcPlatform.GetCoreNumAiv();
  auto blockDim = ascendcPlatform.CalcTschBlockDim(aivNum, aicNum, aivNum);
  context->SetBlockDim(blockDim);
  return blockDim;
}
/*
 *return:
 *0. formerTileLength
 *1. tailTileLength
 *2. formerFinalCalcCount
 *3. tailFinalCalcCount
 *4. finalKernelCaclCount
 */
inline std::tuple<int64_t, int64_t, int64_t, int64_t, int64_t>
splitUb(int64_t split_ub, int64_t formerLength, int64_t tailLength,
        int64_t alignDataSize, int64_t alignedSize, int64_t size) {
  auto formerTileLength = formerLength < split_ub ? formerLength : split_ub;
  auto tailTileLength = tailLength < split_ub ? tailLength : split_ub;
  // 计算finalCalcCount
  auto formerFinalCalcCount = formerLength % formerTileLength;
  formerFinalCalcCount =
      formerFinalCalcCount == 0 ? formerLength : formerFinalCalcCount;
  auto tailFinalCalcCount = tailLength % formerTileLength;
  tailFinalCalcCount =
      tailFinalCalcCount == 0 ? tailLength : tailFinalCalcCount;

  // 计算最后一个核的finalKernelCaclCount
  auto deltaSize = alignedSize - size;
  auto finalKernelCaclCount = tailFinalCalcCount - deltaSize;

  return std::make_tuple(formerTileLength, tailTileLength, formerFinalCalcCount,
                         tailFinalCalcCount, finalKernelCaclCount);
}

inline int64_t
common_tiling_auto_tileLenth(platform_ascendc::PlatformAscendC &ascendcPlatform,
                             int64_t dtypeSize, int64_t alignDataSize,
                             int64_t total_count, int64_t alignedSize) {
  // 计算UB可以容纳多少个数据
  uint64_t ub_size;
  ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ub_size);
  // 计算每个核能容纳多少个操作数
  ub_size /= dtypeSize;
  // ub_size / total_count 然后对齐 alignDataSize
  auto ret = ub_size / total_count / alignDataSize * alignDataSize;
  if (alignedSize < ret) {
    return alignedSize;
  }
  return ret;
}

/*
tiling：TilingData
context: gert::TilingContext
vector_align_size: 对齐大小 默认32, 单位：Byte。
tile_align_size: 一次搬运多少个 vector_align_size 的数据。 默认 8
input_arg_count: 算子的输入个数。 默认:1
output_arg_count: 算子的输出个数。默认:1
auto_tilelength: 是否自动计算 tile_align_size， 如果为true,
tile_align_size参数将失效。默认：true mem_addition： 好需要多少个
tile_align_size 的 空间。一般用来存储Tbuf。 默认：nullptr
target_(dtype/shape)_(input/output)_index:
使用(input/output)的(dtype/shape)来进行切分。 -1 表示不用。默认 input 0
*/
template <typename _T_TILING> class CommonTiling {
public:
  _T_TILING *tiling;
  gert::TilingContext *context = nullptr;
  int64_t data_size;
  uint64_t vector_align_size = 32;
  uint64_t tile_align_size = 8;
  uint64_t input_arg_count = 1;
  uint64_t output_arg_count = 1;
  bool auto_tilelength = true;
  mem_addition_fn mem_addition = nullptr;
  int64_t mem_addition_num = 0;

public:
  CommonTiling() {}
  ~CommonTiling() {}

  void apply(bool use_vector = true, bool use_cube = false) {
    auto ascendcPlatform =
        platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    auto aicNum = ascendcPlatform.GetCoreNumAic();
    auto aivNum = ascendcPlatform.GetCoreNumAiv();
    auto blockDim = ascendcPlatform.CalcTschBlockDim(
        aivNum, use_cube ? aicNum : 0, use_vector ? aivNum : 0);
    context->SetBlockDim(blockDim);

    // 对齐
    auto dtype = context->GetInputDesc(0)->GetDataType();
    auto dtypeSize = GetSizeByDataType(dtype);
    auto alignDataSize = vector_align_size / dtypeSize;
    auto alignCount = (data_size + alignDataSize - 1) / alignDataSize;
    auto alignedSize = alignCount * alignDataSize;

    // 分割
    COMMON_TILING_SPLIT((*tiling));

    // 只有对于搬运耗时较大的运算才使用auto_tilelength
    int64_t formerTileLength = 0;
    int64_t tailTileLength = 0;
    if (auto_tilelength) {
      auto addition = mem_addition_num;
      if (mem_addition) {
        addition += mem_addition(dtype);
      }
      auto total_count = (input_arg_count + output_arg_count) * 2 + addition;
      formerTileLength = common_tiling_auto_tileLenth(
          ascendcPlatform, dtypeSize, alignDataSize, total_count, formerLength);
      tailTileLength = common_tiling_auto_tileLenth(
          ascendcPlatform, dtypeSize, alignDataSize, total_count, tailLength);
    } else {
      auto tileLength = alignDataSize * tile_align_size;
      formerTileLength = formerLength < tileLength ? formerLength : tileLength;
      tailTileLength = tailLength < tileLength ? tailLength : tileLength;
    }

    tiling->set_formerTileLength(formerTileLength);
    tiling->set_tailTileLength(tailTileLength);

    // 最后一次循环
    int64_t formerFinalCalcCount =
        ALIGN_TO((formerLength % formerTileLength), alignDataSize);
    int64_t tailFinalCalcCount =
        ALIGN_TO((tailLength % tailTileLength), alignDataSize);
    if (formerFinalCalcCount == 0) {
      formerFinalCalcCount = formerTileLength;
    }
    if (tailFinalCalcCount == 0) {
      tailFinalCalcCount = tailTileLength;
    }

    tiling->set_formerFinalCalcCount(formerFinalCalcCount);
    tiling->set_tailFinalCalcCount(tailFinalCalcCount);

    // 设置dtype
    tiling->set_dType(dtype);
  }
};

template <typename T>
inline T getDimDefault1(gert::Shape &shape, int64_t index) {
  if (index < 0) {
    index = shape.GetDimNum() + index;
    if (index < 0) {
      return 1;
    }
  }

  if (index > shape.GetDimNum()) {
    return 1;
  }

  return T(shape.GetDim(index));
}

template <typename T0, typename T1, typename T2>
inline std::tuple<bool, T0> getBroadcastDim(T1 dim1, T2 dim2) {
  if (dim1 == dim2) {
    return std::make_tuple(true, dim1);
  }
  if (dim1 == 1) {
    return std::make_tuple(true, dim2);
  }
  if (dim2 == 1) {
    return std::make_tuple(true, dim1);
  }

  return std::make_tuple(false, 0);
}

inline void copyShape(gert::Shape &shape, int64_t *target) {
  for (auto i = 0; i < shape.GetDimNum(); ++i) {
    target[i + 2] = shape.GetDim(i);
  }
  target[0] = shape.GetDimNum();
  target[1] = 0;
}

template <typename T>
inline void print_shape_vector(std::string desc, T shape) {
  std::cout << desc << ": (";
  for (auto i = 1; i <= shape[0]; ++i) {
    std::cout << shape[i + 1] << ", ";
  }
  std::cout << ")" << std::endl;
}

inline void shape_for(gert::Shape &shape,
                      std::function<void(int64_t, int64_t)> func) {
  for (auto i = 0; i < shape.GetDimNum(); ++i) {
    func(i, shape.GetDim(i));
  }
}

inline int64_t getLeafSize(ge::DataType t) {
  switch (t) {
  case ge::DataType::DT_FLOAT16:
    return 200;
  default:
    return 1024;
  }
  return 1024;
}

}; // namespace optiling
#endif